from py2neo import Graph
import os
from dotenv import load_dotenv

load_dotenv()

class Neo4jHandler:
    def __init__(self):
        self.graph = Graph(
            os.getenv("NEO4J_URI"),
            auth=(os.getenv("NEO4J_USER"), os.getenv("NEO4J_PASSWORD"))
        )

    def query_fault_info(self, fault_name):
        """查询故障信息（无fix_code_id时返回None，不抛错）"""
        # 1. 查询故障对应的修复代码ID
        code_query = """
        MATCH (f:FaultType {name: $fault_name})
        RETURN f.fix_code_id AS fix_code_id
        """
        code_result = self.graph.run(code_query, fault_name=fault_name).data()
        fix_code_id = code_result[0]["fix_code_id"] if (code_result and code_result[0]["fix_code_id"]) else None

        # 2. 查询根本原因实体
        entity_query = """
        MATCH (f:FaultType {name: $fault_name})-[r:发生在]->(e:SubEntity)
        RETURN e.name AS root_entity, e.description AS entity_desc
        """
        entity_result = self.graph.run(entity_query, fault_name=fault_name).data()
        if not entity_result:
            # 允许无实体信息，后续由Qwen补充
            root_entity = {"root_entity": "未知实体", "entity_desc": "未明确根因实体"}
        else:
            root_entity = entity_result[0]

        # 3. 查询上游依赖关系
        dependency_query = """
        MATCH (f:FaultType {name: $fault_name})-[r:发生在]->(e:SubEntity)
        MATCH (e)-[:依赖于*1..]->(upstream:SubEntity)
        RETURN upstream.name AS upstream_entity, collect(distinct r.description) AS dependency_desc
        """
        dependencies = self.graph.run(dependency_query, fault_name=fault_name).data()
        for dep in dependencies:
            dep["dependency_desc"] = "; ".join(dep["dependency_desc"]) if dep["dependency_desc"] else "无描述"

        # 4. 查询知识图谱中的解决方案
        solution_query = """
        MATCH (f:FaultType {name: $fault_name})<-[r:用于解决]-(sol:Solution)
        RETURN sol.content AS solution
        """
        solution_result = self.graph.run(solution_query, fault_name=fault_name).data()
        solutions = [item["solution"] for item in solution_result] if solution_result else ["无预定义解决方案"]

        return {
            "fault_name": fault_name,
            "fix_code_id": fix_code_id,  
            "root_cause": {
                "entity": root_entity["root_entity"],
                "description": root_entity["entity_desc"]
            },
            "dependencies": dependencies,
            "solutions_from_kg": solutions
        }

    def query_upstream_fault_info(self, upstream_entity):
        """保持不变，确保上游故障信息能正常获取"""
        fault_query = """
        MATCH (e:SubEntity {name: $upstream_entity})<-[r:发生在]-(f:FaultType)
        RETURN f.name AS fault_name
        """
        fault_result = self.graph.run(fault_query, upstream_entity=upstream_entity).data()
        if not fault_result:
            return []
        upstream_faults = []
        for item in fault_result:
            fault_name = item["fault_name"]
            fault_info = self.query_fault_info(fault_name)
            upstream_faults.append(fault_info)
        return upstream_faults